"""
手写体识别：本地文件
"""
import gzip
import json
import random

import numpy as np
import paddle
from paddle.nn import Conv2D,MaxPool2D,Linear
import paddle.nn.functional as F


def load_data(mode='train'):
    datafile = './work/mnist.json.gz'
    print('loading mnist dataset from {}'.format(datafile))
    # 加载json数据文件
    data = json.load(gzip.open(datafile))
    print('mnist dataset load done')
    # 读取到的数据区分训练集，验证集，测试集
    train_set,val_set,eval_set = data

    if mode == 'train':
        # 观察训练集数据
        imgs,labels = train_set[0],train_set[1]
        # print("训练数据集数量：",len(imgs))
    elif mode == 'valid':
        # 观察验证集数量
        imgs,labels = val_set[0],val_set[1]
        # print("验证数据集数量：",len(imgs))
    elif mode == 'eval':
        # 观察测试集数量
        imgs,labels = eval_set[0],eval_set[1]
        # print('测试数据集数量：',len(imgs))
        # print(len(imgs[0]))
    else:
        raise Exception("mode can only be one of ['train','valid','eval']")
    print("训练数据集数量：",len(imgs))

    assert len(imgs) == len(labels),"length of train_imgs({}) should be the same as train_labels({})".format(len(imgs),len(labels))
    # 获得数据集长度
    imgs_length = len(imgs)
    # 定义数据集每个数据的序号，根据序号读取数据
    index_list = list(range(imgs_length))
    # 读入数据时用到的批次大小
    BATCHSIZE = 100

    # 定义数据生成器，返回批次数据
    def data_generator():
        if mode == 'train':
            # 随机打乱训练数据的索引序号
            random.shuffle(index_list)
        imgs_list = []
        labels_list = []
        for i in index_list:
            # 将数据处理成希望的类型
            img = np.array(imgs[i]).astype('float32').reshape(1,28,28)
            label = np.array(labels[i]).astype('float32')
            imgs_list.append(img)
            labels_list.append(label)
            if len(imgs_list) == BATCHSIZE:
                # 获得一个batchsize的数据，并返回
                yield np.array(imgs_list), np.array(labels_list)
                # 清空数据读取列表
                imgs_list = []
                labels_list = []

        # 如果剩余的数据的数目小于BATCHSIZE
        # 则剩余数据一起构成一个大小为len(imgs_list)的mini-batch
        if len(imgs_list) > 0:
            yield np.array(imgs_list),np.array(labels_list)

    return data_generator

    # # 声明数据读取函数，从训练集中读取数据
    # train_loader = data_generator
    # # 以迭代的形式读取数据
    # for batch_id,data in enumerate(train_loader()):
    #     image_data,label_data = data
    #     if batch_id == 0:
    #         print("打印第一个batch数据的维度：")
    #         print("图像维度：{}，标签维度：{}".format(image_data.shape,label_data.shape))
    #     break

class MNIST(paddle.nn.Layer):
    def __init__(self):
        super(MNIST,self).__init__()

        # 定义卷积层，输出特征通道out_channels设置为20，卷积核的大小kernel_size为5，卷积步长stride=1, padding=2
        self.conv1 = Conv2D(in_channels=1,out_channels=20,kernel_size=5,stride=1,padding=2)
        # 定义池化层，池化核的大小kernel_size为2，池化步长为2
        self.max_pool1 = MaxPool2D(kernel_size=2,stride=2)
        # 定义卷积层，输出特征通道out_channels设置为20，卷积核的大小kernel_size为5，卷积步长stride=1,padding=2
        self.conv2 = Conv2D(in_channels=20,out_channels=20,kernel_size=5,stride=1,padding=2)
        # 定义池化层，池化核的大小kernel_size为2，池化步长为2
        self.max_pool2 = MaxPool2D(kernel_size=2,stride=2)
        # 定义一层全连接层,输出维度是1
        self.fc = Linear(in_features=980,out_features=10)

    def forward(self,inputs):
        """
        定义网络的前向计算，卷积后紧接着使用池化层，最后使用全连接层计算最终输出
        卷积层激活函数使用Relu,全连接层不适用激活函数
        :param inputs:
        :return:
        """
        x = self.conv1(inputs)
        x = F.relu(x)
        x = self.max_pool1(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.max_pool2(x)
        x = paddle.reshape(x,[x.shape[0],980])
        x = self.fc(x)
        return x

# model = MNIST()
# params_info = paddle.summary(model,(1,1,28,28))
# print(params_info)

def train(model):
    # model = MNIST()
    # 开启GPU
    use_gpu = True
    paddle.device.set_device('gpu:0') if use_gpu else paddle.device.set_device('cpu')

    model.train()

    learning_rate = 0.001
    # 调用加载数据的函数
    train_loader = load_data('train')

    opt = paddle.optimizer.SGD(learning_rate=learning_rate,parameters=model.parameters())
    EPOCH_NUM = 10
    # MNIST图像高和宽
    # IMG_ROWS,IMG_COLS = 28,28
    loss_list = []
    for epoch_id in range(EPOCH_NUM):
        for batch_id,data in enumerate(train_loader()):
            # 准备数据，变的更加简洁
            images,labels = data
            images = paddle.to_tensor(images)
            labels = paddle.to_tensor(labels)

            # 前向计算的过程
            predits = model(images)

            # 计算损失，取一个批次样本损失的平均值
            loss = F.cross_entropy(predits,labels)
            avg_loss = paddle.mean(loss)

            # 每训练了200批次的数据，打印下当前loss的情况
            if batch_id % 200 ==0:
                loss_list.append(avg_loss.numpy())
                print("epoch: {},batch: {},loss is {}".format(epoch_id,batch_id,avg_loss.numpy()))

            #后向传播，更新参数的过程
            avg_loss.backward()
            # 最小化loss,更新参数
            opt.step()
            # 清除梯度
            opt.clear_grad()

    # 保存模型
    paddle.save(model.state_dict(),'./mnist5.pdparams')
    return loss_list

# # 创建模型
# model = MNIST()
# # 启动训练过程
# loss_list = train(model)
print(load_data('train'))